from langchain_community.vectorstores import Qdrant
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, PyPDFLoader
import os
import json
import warnings
from openai import OpenAI
warnings.filterwarnings('ignore')

# 读取API Key
with open(r'd:\code\vscode\.Python\taidi\code_tadi\ERNIE_config.json') as f:
    key = f.read()
    api_key = json.loads(key)

# 设置通义千问API Key
DASHSCOPE_API_KEY = api_key['DASHSCOPE_API_KEY']  # 阿里通义千问模型DASHSCOPE_API_KEY
os.environ["DASHSCOPE_API_KEY"] = DASHSCOPE_API_KEY  # 设置环境变量

# 创建OpenAI客户端（通义千问兼容模式）
client = OpenAI(
    api_key=DASHSCOPE_API_KEY,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 加载PDF文档
def load_documents(pdf_path):
    """加载PDF文档并进行分块处理"""
    try:
        # 检查文件是否存在
        if not os.path.exists(pdf_path):
            print(f"错误: 文件不存在 - {pdf_path}")
            return []
            
        loader = PyPDFLoader(pdf_path)  # 读取数据
        documents = loader.load()  # 将数据转为指定格式
        
        # 文本分割
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)  # 指定切分时每份文件的token数
        chunked_documents = text_splitter.split_documents(documents)  # 将文本切分成指定长度大小
        
        if not chunked_documents:
            print(f"警告: 文档加载成功但未提取到内容 - {pdf_path}")
            
        return chunked_documents
    except Exception as e:
        print(f"文档加载错误: {e}")
        return []

# 创建向量数据库
def create_vectorstore(documents):
    """创建向量数据库"""
    try:
        # 使用OpenAI兼容的嵌入模型
        from langchain_openai import OpenAIEmbeddings
        
        # 配置嵌入模型（通义千问兼容模式）
        try:
            embedding = OpenAIEmbeddings(
                model="text-embedding-v1",  # 使用通义千问支持的嵌入模型
                openai_api_key=DASHSCOPE_API_KEY,
                openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
            )
            print(f"嵌入模型初始化成功: text-embedding-v1")
        except Exception as e:
            print(f"嵌入模型初始化错误: {e}")
            return None
        
        # 利用Embedding将文本向量化，并存储在向量数据库中
        vectorstore = Qdrant.from_documents(
            documents=documents,  # 已分块的文档
            embedding=embedding,  # 使用嵌入模型
            location=":memory:",  # 存储在内存中
            collection_name="my_documents"
        )  # 指定collection_name
        
        return vectorstore
    except Exception as e:
        print(f"向量数据库创建错误: {e}")
        return None

# 创建RAG链
def create_rag_chain(vectorstore):
    """创建RAG检索链"""
    # 创建检索器
    retriever = vectorstore.as_retriever()
    
    # 配置上下文（context）及用户输入（question）
    setup_and_retrieval = RunnableParallel(
        {"context": retriever, "question": RunnablePassthrough()}
    )
    
    # 创建提示模板
    template = """根据以下提供的上下文信息，回答问题。只使用上下文中的信息，如果上下文中没有相关信息，请说明无法回答。
    
    上下文信息:
    {context}
    
    问题: {question}
    
    回答:"""
    
    prompt = ChatPromptTemplate.from_template(template)
    
    # 创建通义千问模型（使用langchain_openai的ChatOpenAI）
    from langchain_openai import ChatOpenAI
    
    model = ChatOpenAI(
        model="qwen-plus",  # 使用通义千问Plus模型
        temperature=0.7,  # 控制生成文本的随机性
        openai_api_key=DASHSCOPE_API_KEY,
        openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )
    
    # 输出解析器
    output_parser = StrOutputParser()
    
    # 构建完整的链（Chain）
    chain = setup_and_retrieval | prompt | model | output_parser
    
    return chain

# 主RAG类
class QwenRAG:
    """通义千问RAG系统"""
    def __init__(self, pdf_path):
        print(f"正在加载文档: {pdf_path}")
        self.documents = load_documents(pdf_path)
        if self.documents:
            print(f"文档加载成功，共 {len(self.documents)} 个文本块")
            print("正在创建向量数据库...")
            self.vectorstore = create_vectorstore(self.documents)
            if self.vectorstore:
                print("向量数据库创建成功")
                print("正在初始化RAG链...")
                self.chain = create_rag_chain(self.vectorstore)
                print("RAG系统初始化完成")
            else:
                print("向量数据库创建失败")
                self.chain = None
        else:
            print("文档加载失败，无法初始化RAG系统")
            self.vectorstore = None
            self.chain = None
    
    def ask(self, question):
        """向RAG系统提问"""
        if self.chain:
            try:
                return self.chain.invoke(question)
            except Exception as e:
                return f"查询处理错误: {e}"
        else:
            return "RAG系统未正确初始化，请检查文档加载和向量数据库创建过程。"

# 示例使用
if __name__ == "__main__":
    # 初始化RAG系统
    # 使用实际存在的PDF文件路径
    pdf_path = 'd:\\code\\vscode\\.Python\\taidi\\subject\\C题-竞赛智能客服机器人.pdf'  # 使用绝对路径
    rag = QwenRAG(pdf_path)
    
    # 测试问题
    test_questions = [
        "未来校园智能应用专项赛的主要内容是什么？",
        "这个比赛的参赛对象有哪些？",
        "比赛的评分标准是什么？"
    ]
    
    for question in test_questions:
        print(f"问题: {question}")
        answer = rag.ask(question)
        print(f"回答: {answer}")
        print("-" * 50)